#!/usr/bin/env python3
"""Run faster-whisper on V2 voiceover to get word-level aligned captions"""

import os, json, sys, subprocess, re

BASE = r'E:\集群文件夹\factory_os\short_video_real_data_pipeline\phase4b_90s_formal_sample\v22_kimi_claude_loop\full_story_douyin_video'
AUDIO = os.path.join(BASE, '02_voiceover', 'v2', 'voiceover_master_v2.wav')
SUB_DIR = os.path.join(BASE, '04_subtitles', 'v2')
os.makedirs(SUB_DIR, exist_ok=True)

print(f"Using audio: {AUDIO}")
print(f"Output dir: {SUB_DIR}")

# Run faster-whisper
from faster_whisper import WhisperModel

# Use cached base model (large-v3 incomplete, no internet access for download)
model = WhisperModel("base", device="cuda", compute_type="float16", local_files_only=True)
print("Using GPU (CUDA) with base model")

print("\nTranscribing V2 voiceover...")
segments, info = model.transcribe(AUDIO, language="zh", beam_size=5, word_timestamps=True)

print(f"Detected language: {info.language} (p={info.language_probability:.2f})")

# Collect word-level results
words_all = []
segments_list = []
for seg in segments:
    seg_data = {
        "start": round(seg.start, 2),
        "end": round(seg.end, 2),
        "text": seg.text.strip(),
        "words": []
    }
    if seg.words:
        for w in seg.words:
            wd = {
                "word": w.word.strip(),
                "start": round(w.start, 2),
                "end": round(w.end, 2),
                "probability": round(w.probability, 3) if w.probability else 0
            }
            seg_data["words"].append(wd)
            words_all.append(wd)
    segments_list.append(seg_data)
    print(f"  [{seg.start:.1f}s-{seg.end:.1f}s] {seg.text.strip()[:60]}")

# Save word-level JSON
captions_words = {
    "total_duration": 92.6,
    "model": "faster-whisper-medium",
    "segments": segments_list,
    "word_count": len(words_all)
}

words_path = os.path.join(SUB_DIR, 'captions_words_v2.json')
with open(words_path, 'w', encoding='utf-8') as f:
    json.dump(captions_words, f, ensure_ascii=False, indent=2)
print(f"\nWord-level captions saved: {words_path}")

# Now generate SRT from word-level data
def format_srt_time(seconds):
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}".replace('.', ',')

# Build SRT - group words into subtitle lines
# Max 2 lines per screen, max 13 Chinese chars per line
srt_lines = []
sub_idx = 1

# Keywords to highlight
HIGHLIGHTS = ["录屏失败", "AI 跑偏", "真实日志", "Kimi", "69", "92", "76", "88",
              "regression", "回滚", "自动修复", "稳定交付", "HTTP", "200"]

def should_highlight(word):
    w = word.strip().rstrip('.,!?；。，！？、')
    for h in HIGHLIGHTS:
        if h.lower() in w.lower() or w.lower() in h.lower():
            return True
    return False

# Group words into subtitle groups (each group ~2 lines, ~13 chars per line)
subtitle_groups = []
current_group = []
current_text = ""
max_chars_per_group = 26  # 2 lines * 13 chars

for seg in segments_list:
    for w in seg["words"]:
        word_text = w["word"].strip()
        if not word_text:
            continue

        # Check if adding this word would exceed max chars
        if len(current_text + word_text) > max_chars_per_group and current_group:
            # Save current group
            subtitle_groups.append({
                "words": current_group,
                "text": current_text.strip()
            })
            current_group = []
            current_text = ""

        current_group.append(w)
        current_text += word_text

# Don't forget last group
if current_group:
    subtitle_groups.append({
        "words": current_group,
        "text": current_text.strip()
    })

# Generate SRT with proper timing
srt_path = os.path.join(SUB_DIR, 'captions_final_v2.srt')
ass_path = os.path.join(SUB_DIR, 'captions_final_v2.ass')

# SRT
with open(srt_path, 'w', encoding='utf-8') as f:
    for i, group in enumerate(subtitle_groups):
        words = group["words"]
        start = words[0]["start"]
        end = words[-1]["end"]

        # Minimum display time 0.5s
        if end - start < 0.5:
            end = start + 0.5

        f.write(f"{i+1}\n")
        f.write(f"{format_srt_time(start)} --> {format_srt_time(end)}\n")

        # Split into 2 lines if too long
        text = group["text"]
        if len(text) > 13:
            # Try to split at natural boundary
            split_pos = len(text) // 2
            # Find nearest space or punctuation
            for j in range(len(text) // 2, 0, -1):
                if text[j] in ' ，。！？、':
                    split_pos = j + 1
                    break
            line1 = text[:split_pos].strip()
            line2 = text[split_pos:].strip()
            f.write(f"{line1}\n{line2}\n\n")
        else:
            f.write(f"{text}\n\n")

print(f"SRT saved: {srt_path}")

# ASS subtitle with keyword highlighting
# ASS header
ass_header = """[Script Info]
ScriptType: v4.00+
PlayResX: 1080
PlayResY: 1920
ScaledBorderAndShadow: yes

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,Microsoft YaHei,28,&H00FFFFFF,&H0000FF00,&H00000000,&H80000000,1,0,0,0,100,100,0,0,1,2,0,2,50,50,80,134
Style: Keyword,Microsoft YaHei,30,&H0000FFDD,&H0000FF00,&H00000000,&H80000000,1,0,0,0,100,100,0,0,1,2,0,2,50,50,80,134

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""

with open(ass_path, 'w', encoding='utf-8') as f:
    f.write(ass_header)

    for i, group in enumerate(subtitle_groups):
        words = group["words"]
        start = words[0]["start"]
        end = words[-1]["end"]

        if end - start < 0.5:
            end = start + 0.5

        start_str = format_srt_time(start).replace(',', '.')
        end_str = format_srt_time(end).replace(',', '.')

        # Build ASS line with keyword highlights
        ass_text_parts = []
        for w in words:
            word_text = w["word"].strip()
            if not word_text:
                continue
            if should_highlight(word_text):
                ass_text_parts.append(f"{{\\c&H00FFDD&\\b1}}{word_text}{{\\c&HFFFFFF&\\b0}}")
            else:
                ass_text_parts.append(word_text)

        ass_line = " ".join(ass_text_parts)

        # Split into 2 lines if too long
        if len(group["text"]) > 13:
            # Simple split at midpoint
            raw = group["text"]
            mid = len(raw) // 2
            for j in range(mid, 0, -1):
                if raw[j] in ' ，。！？、':
                    mid = j
                    break
            # Split the ass text too
            char_count = 0
            split_idx = 0
            for pi, part in enumerate(ass_text_parts):
                # Count chars in this part (strip ASS tags)
                clean = re.sub(r'\{[^}]*\}', '', part)
                if char_count + len(clean) > mid:
                    split_idx = pi
                    break
                char_count += len(clean)

            if split_idx > 0:
                line1_parts = ass_text_parts[:split_idx]
                line2_parts = ass_text_parts[split_idx:]
                line1_text = "".join(line1_parts)
                line2_text = "".join(line2_parts)
                combined = f"{line1_text}\\N{line2_text}"
            else:
                combined = "".join(ass_text_parts)
        else:
            combined = "".join(ass_text_parts)

        # Write dialogue line with Default style
        f.write(f"Dialogue: 0,{start_str},{end_str},Default,,0,0,0,,{combined}\n")

print(f"ASS saved: {ass_path}")
print(f"\nTotal subtitle groups: {len(subtitle_groups)}")
print(f"Total words: {len(words_all)}")
print("\nWord-level caption generation complete!")